#!/usr/bin/perl -w
#---------------------------------------------------------------------------
#@COPYRIGHT :
#             Copyright 1996, John G. Sled, 
#             McConnell Brain Imaging Centre,
#             Montreal Neurological Institute, McGill University.
#             Permission to use, copy, modify, and distribute this
#             software and its documentation for any purpose and without
#             fee is hereby granted, provided that the above copyright
#             notice appear in all copies.  The author and McGill University
#             make no representations about the suitability of this
#             software for any purpose.  It is provided "as is" without
#             express or implied warranty.
#---------------------------------------------------------------------------- 
#$RCSfile: nu_estimate.in,v $
#$Revision: 1.1 $
#$Author: bert $
#$Date: 2003/04/16 14:29:34 $
#$State: Exp $
#---------------------------------------------------------------------------
# ------------------------------ MNI Header ----------------------------------
#@NAME       : nu_estimate or nu_correct
#@INPUT      : 
#@OUTPUT     : 
#@RETURNS    : 
#@DESCRIPTION: estimate intensity non-uniformity artifacts from given volume
#@            or
#@             estimate and correct intensity non-uniformity artifact  
#@METHOD     : 
#@GLOBALS    : 
#@CALLS      : 
#@CREATED    : August 7, 1996      J.G.Sled  
#@MODIFIED   : 
#  $Id: nu_estimate.in,v 1.1 2003/04/16 14:29:34 bert Exp $
#-----------------------------------------------------------------------------

use MNI::Startup qw(nocputimes);
use MNI::Spawn;
use MNI::FileUtilities qw(search_directories);
use MNI::PathUtilities qw(replace_ext replace_dir);
use MNI::MincUtilities qw(:history);
use Getopt::Tabular;

# Start main program ------------------------------------------------
&Initialize;

my (@Outputs, @Inputs);

# run nu_estimate
$nu_estimate = ($estimation_type eq 'wm') ? 
   'nu_estimate_wm' : 'nu_estimate_np_and_em';

if($program_type eq 'estimate') {
      Spawn("$nu_estimate $estimation_options @Inputs @Outputs");
}
else { #run nu_correct
   @Imps = &replace_ext('imp', @Outputs);
   (defined $mapping_dir) && (@Imps = &replace_dir($mapping_dir,@Imps));

   Spawn("$nu_estimate $estimation_options @Inputs @Imps");
   for($i = 0; $i < $#Inputs+1; $i++) {
      Spawn(['nu_evaluate', '-shrink', 3,
             $Inputs[$i], '-mapping', $Imps[$i], $Outputs[$i]]);
      &construct_history($Inputs[$i], $Outputs[$i], "$0 @ARGV");
   }
}
#  End of Main Program
#-----------------------------------------------------------------------------

# ------------------------------ MNI Header ----------------------------------
#@NAME       : &construct_history
#@INPUT      : $_[0]  name of input volume 
#              $_[1]  name of output volume to get new history
#              $_[2]  string to append to history
#@OUTPUT     : 
#@RETURNS    : 
#@DESCRIPTION: 
#@METHOD     : 
#@GLOBALS    : 
#@CALLS      : get_history, put_history
#@CREATED    : 
#@MODIFIED   : 
#-----------------------------------------------------------------------------
sub construct_history
{
   my (@history);
   @history  = &get_history($_[0]);
   push(@history, $_[2]);
   &put_history($_[1], @history);
}


# ------------------------------ MNI Header ----------------------------------
#@NAME       : &SetupArgTables
#@INPUT      : none
#@OUTPUT     : none
#@RETURNS    : References to the two option tables:
#                @pref_args,
#                @protocol_args
#                @correct_args
#                @np_em_args
#                @np_args
#@DESCRIPTION: Defines the tables of command line (and config file) 
#              options that we pass to ParseArgs.
#@METHOD     : 
#@GLOBALS    : makes references to many globals (almost all of 'em in fact)
#              even though most of them won't have been defined when
#              this is called
#@CALLS      : 
#@CREATED    : 
#@MODIFIED   : 
#-----------------------------------------------------------------------------
sub SetupArgTables
{
   my ($N3_only_flag) = @_;
   my (@pref_args, @protocol_args);

   sub print_version
   {
      print "Program $ProgramName, built from:\n$LongVersion\n";
      exit;
   }

   # Preferences -- these shouldn't affect the results

   @pref_args = 
      (["General options", "section"],
       ["-verbose|-quiet", "boolean", 0, \$Verbose, 
	"be noisy [default; opposite is -quiet]" ],
       ["-execute|-noexecute", "boolean", 0, \$Execute, 
	"actually execute planned commands [default]"],
       ["-clobber|-noclobber", "boolean", 0, \$Clobber,
	"overwrite output files [default: -noclobber]"],
#       ["-compress|-nocompress", "boolean", 0, \$Compress,
#	"compress output files [default: -nocompress]"],
       ["-tmpdir", "string", 1, \$TmpDir,
	"temporary working directory"],
       ["-keeptmp|-nokeeptmp", "boolean", 0, \$KeepTmp,
	"don't delete temporary files [default: -nokeeptmp]"],
       ["-version", "call", undef, \&print_version,
        "print version and quit"]);
#       ["-notify", "string", 1, \$Notify,
#	"set the user to notify (send email to) on failure "],
#       ["-nonotify", "constant", 0, \$Notify,
#	"disable email notification (default)"]);

   # Protocol (data-specific) options 

   @protocol_args = 
      (["Protocol options","section"],
       (
        ($N3_only_flag) ? () :
        (["-non_parametric", "const", 'np', \$estimation_type,
          "use non-parametric (N3) correction method (default)"],
         ["-white_matter", "const", 'wm', \$estimation_type,
          "use white matter based correction method"],
         ["-expectation_maximization", "const", 'em',
          \$estimation_type,  
          "use expectation maximization correction method"])
        ),
       ["-mask", "volume", 1, \$user_options{'mask'},
	"specify region for processing", "<mask.mnc>"],
       ["-distance", "float", 1, \$user_options{'distance'}, 
        "characteristic distance over which field varies in mm (default 200)",
        "<value>"],
       ["-iterations", "integer", 1, \$user_options{'iterations'},
	"maximum number of iterations to run", "<number>"],
       ["-options", "string", 1, \$estimation_options,
        "method specific options (quoted if more than one)",<'options ...'>]); 

   @correct_args =  (["-mapping_dir", "string", 1, \$mapping_dir, 
                      "specifiy directory to put imp files in","<directory>"]);
   
   @np_em_args = ((($N3_only_flag) ? 
                   (["N3 specific options", "section"]) : 
                   (["N3 and EM specific options", "section"])),
                  ["-stop", "float", 1, \$user_options{'stop'},
                   "CV of change in field estimate below which"
                   ." iteration stops (suggest 0.01 to 0.0001)",
                   "<threshold>" ],
                  ["-shrink", "float", 1, \$user_options{'shrink'},
                   "specify a smaller workspace; the"
                   ." sampling in each dimension is reduced to that of"
                   ." the finest"
                   ." sampling divided by factor or the current sampling"
                   ." which ever"
                   ." is less. (suggest 2 or 3)", "<factor>"],
                  ["-normalize_field", "const", 1, 
                   \$user_options{'normalize_field'}, 
                   "normalize final correction field to have mean 1.0"]);

   @np_args = ((($N3_only_flag) ? () :
                (["N3 specific options","section"])),
               ["-V0.9", "const", '0.9', \$version_number,
                "use default options of original (beta) release"],
               ["-V1.0", "const", '1.0', \$version_number,
                "use improved options of release 1.0.  This is equivalent to "
                ."-fwhm 0.15 -stop 0.001 -iterations 50 -shrink "
                ."4 -distance 200.  (default)"],
#               ["-newest", "const", '1.0', \$version_number,
#                "use newest version of software."],
               ["-fwhm", "float", 1, \$user_options{'sharpen'}, 
                "width of deconvolution kernal used to"
                ." sharpen histogram.  Larger values give faster convergence"
                ." while smaller values give greater accuracy. "
                ." (suggest 0.05 to 0.5, default is 0.15)","<width>"]);

   (\@pref_args, \@protocol_args, \@correct_args, \@np_em_args, \@np_args);
}

# ------------------------------ MNI Header ----------------------------------
#@NAME       : &SetHelp
#@INPUT      : $program_type = 'estimate' or 'correct'
#@OUTPUT     : none
#@RETURNS    : nothing
#@DESCRIPTION: Sets the $Help and $Usage globals, and registers them
#              with ParseArgs so that user gets useful error and help
#              messages.
#@METHOD     : 
#@GLOBALS    : $Help, $Usage
#@CALLS      : 
#@CREATED    : 
#@MODIFIED   : 
#-----------------------------------------------------------------------------
sub SetHelp
{
   my($program_type, $N3_only_flag) = @_;

   $Version = '1.10';
   $LongVersion = 'Package MNI N3, version 1.10, compiled by nicks@minerva (x86_64-unknown-linux-gnu) on 2005-11-15 at 20:37:35';

   # common help message
   if($N3_only_flag) {
      $CommonHelp = <<COMMONHELP;

       The N3 method [1] should work with any MR volume including raw
       (non-stereotaxic) data.  The performance of this method can be
       enhanced by supplying a mask for the region of interest.

   References: 
   [1] J.G. Sled, A.P. Zijdenbos and A.C. Evans, "A non-parametric method
       for automatic correction of intensity non-uniformity in MRI data",
       in "IEEE Transactions on Medical Imaging", vol. 17, n. 1,
       pp. 87-97, 1998 

COMMONHELP

      if($program_type eq 'estimate') {
      $Usage = <<ESTIMATEUSAGE;
$ProgramName, version $Version

Usage: $ProgramName [-help] [options] input.mnc output.imp

ESTIMATEUSAGE

      $Help = <<ESTIMATEHELP . $CommonHelp;
       $ProgramName is a simplified interface to the script 
       nu_correct_np_and_em for estimating intensity non-uniformity in
       MR volumes.  The script nu_evaluate should be used to apply the
       resulting intensity mapping(s) to correct the source volume(s).
ESTIMATEHELP
      }
      else {
      $Usage = <<CORRECTUSAGE;
$ProgramName, version $Version

Usage: $ProgramName [-help] [options] in.mnc out.mnc

CORRECTUSAGE
      $Help = <<CORRECTHELP . $CommonHelp;
       $ProgramName is a simplified interface to the scripts nu_estimate
       and nu_evaluate for correcting for intensity non-uniformity in
       MR volumes.
CORRECTHELP

      }
   }
   else {
   $CommonHelp = <<COMMONHELP;

       The N3 method [1], default, should work with any MR volume
       including raw (non-stereotaxic) data.  The performance of this
       method can be enhanced by supplying a mask for the region of
       interest.  Correction of multiple modalities at once is not
       supported.

       The white matter method is a slightly modified verison of the
       algorithm previously implemented in MIDAS [2]. The input files are
       taken as multiple modalities of the same volume, which are used
       simultaneously in the classification stage.  These volumes must live
       in Talairach space.

       The expectation maximization methed is a modified implementation of
       the algorithm described in [3].  Like the white matter method, volumes
       must reside in Talairach space, however correction of multiple
       modalities is not supported.

   References: 
   [1] J.G. Sled, A.P. Zijdenbos and A.C. Evans, "A non-parametric method
       for automatic correction of intensity non-uniformity in MRI data",
       in "IEEE Transactions on Medical Imaging", vol. 17, n. 1,
       pp. 87-97, 1998 

   [2] A.P. Zijdenbos, B.M. Dawant and R.A. Margolin, "Intensity Correction
       and its Effect on Measurement Variability in the Computer-Aided
       Analysis of {MRI}", in "Proceedings of the 9th International
       Symposium and Exhibition on Computer Assisted Radiology (CAR)",
       Berlin, Germany, June 1995, pp. 216-221.

   [3] W.M. Wells III, W.E.L. Grimson, R. Kikinis and F. A. Jolesz,
       "Adaptive Segmentation of MRI Data", in "IEEE Transactions on
       Medical Imaging", volume 15, number 4, 1996, pp. 429-442

COMMONHELP

      if($program_type eq 'estimate') {
      $Usage = <<ESTIMATEUSAGE;
$ProgramName, version $Version

Usage: $ProgramName [-help] [options] input.mnc output.imp
     
  or   $ProgramName -white_matter in.mnc [in2.mnc ...] out.imp [out2.imp ...]
ESTIMATEUSAGE

      $Help = <<ESTIMATEHELP . $CommonHelp;
       $ProgramName is a simplified interface to the scripts 
       nu_correct_np_and_em and nu_correct_wm for estimating intensity
       non-uniformity in MR volumes.  The script nu_evaluate should be
       used to apply the resulting intensity mapping(s) to correct the
       source volume(s).
ESTIMATEHELP
      }
      else {
      $Usage = <<CORRECTUSAGE;
$ProgramName, version $Version

Usage: $ProgramName [-help] [options] in.mnc out.mnc

  or   $ProgramName [options] in1.mnc in2.mnc ... out1.mnc out2.mnc ...

CORRECTUSAGE
      $Help = <<CORRECTHELP . $CommonHelp;
       $ProgramName is a simplified interface to the scripts nu_estimate
       and nu_evaluate for correcting for intensity non-uniformity in
       MR volumes.
CORRECTHELP

      }
   }
   Getopt::Tabular::SetHelp ($Help, $Usage);       
}


# ------------------------------ MNI Header ----------------------------------
#@NAME       : &search_for_programs
#@INPUT      : \@programs, $path
#@OUTPUT     : 
#@RETURNS    : 1 if all programs are found, 0 otherwise
#@DESCRIPTION: checks for given list of programs in given directories
#@METHOD     : 
#@GLOBALS    : 
#@CALLS      : search_directories
#@CREATED    : 
#@MODIFIED   : 
#-----------------------------------------------------------------------------
sub search_for_programs 
{
   my ($programs, $path) = @_;
   my (@path_list) = split(':',$path);
   my ($program);
   
   foreach $program (@$programs) {
      if(!search_directories($program, \@path_list, "-x")) {
         return 0;
      }
   }
   return 1;
}

# ------------------------------ MNI Header ----------------------------------
#@NAME       : &Initialize
#@INPUT      : 
#@OUTPUT     : 
#@RETURNS    : 
#@DESCRIPTION: Sets global variables, parses command line.  Dies on
#              any error.
#@METHOD     : 
#@GLOBALS    : preferences: $Verbose, $Execute, $Clobber, $Debug, $KeepTmp
#              
#@CALLS      : &SetSpawnOptions
#              &SetupArgTables
#              &ParseArgs::Parse
#@CREATED    : 
#@MODIFIED   : 
#-----------------------------------------------------------------------------
sub Initialize
{
   my ($pref_tbl, $protocol_tbl, @reducedARGV, @programs);

   # Set defaults for the global variables.  These can be overridden by 
   # the command line.
   $Verbose      = 0;

   # determine what functionality is available
   $N3_only_flag = ! &search_for_programs(['nu_estimate_np_and_em',
                                           'nu_estimate_wm', 'nu_evaluate',
                                           'class_statistics','estimate',
                                           'lgmask'],$ENV{'PATH'});

   # determine which program is being asked for
   $program_type = ($ProgramName =~ /nu_estimate/i) ? 'estimate' : 'correct';
   &SetHelp($program_type, $N3_only_flag);

   # protocol options
   undef $user_options{'iterations'};
   undef $user_options{'mask'};
   undef $user_options{'stop'};
   undef $user_options{'shrink'};
   $user_options{'distance'} = 200;
   $user_options{'sharpen'} = 0.15;
   $version_number = '1.0'; 
   undef $mapping_dir; 

   @option_names = ('iterations', 'stop', 'shrink');
   %defaults = ('np:iterations:0.9' => '10 20',         
                'np:iterations:1.0' => '50',
                'np:stop:0.9'       => '0.001 0.005',
                'np:stop:1.0'       => '0.001',
                'np:shrink:0.9'     => '3',             
                'np:shrink:1.0'     => '4',
                
                'em:iterations'     => '10',
                'em:stop'           => '0.001',
                'em:shrink'         => '3');


   # more global variables
   $estimation_type = 'np';   #default
   $estimation_options = '';

   # Parse command line arguments
   ($pref_tbl, $protocol_tbl, $correct_tbl, $np_em_tbl, $np_tbl ) = 
      &SetupArgTables($N3_only_flag);
   @all_args = ($program_type eq 'estimate') ?
      (@$pref_tbl, @$protocol_tbl, @$np_em_tbl, @$np_tbl) :
         (@$pref_tbl, @$protocol_tbl, @$correct_tbl, @$np_em_tbl, @$np_tbl);
   &Getopt::Tabular::AddPatternType("volume", ".+(\\.mnc|\\.mnc\\.gz|"
                       ."\\.mnc\\.Z|\\.mnc\\.z)\$", "minc volume");
   GetOptions(\@all_args, \@ARGV, \@reducedARGV)
         || exit 1;

   if(@reducedARGV == 0)
   {
      print STDERR $Usage;
      exit 2;
   }
   else # Check source arguments
   {
      my($nArgs) = $#reducedARGV + 1;
      if (!$nArgs || ($nArgs % 2)) {
         die "Please supply an even number of arguments.\n";
      }
      @Inputs = splice(@reducedARGV, 0, $nArgs/2);
      @Outputs = @reducedARGV;
   }
   
   if(!$Clobber) {
      foreach $output (@Outputs) {
         (-e $output || -e "$output.gz" ) &&
            (die("Clobber option not given.  Cannot overwrite file:"
                 ." $output\n")); 
      }
   }
   # Set global variables for calling various programs
   MNI::Spawn::SetOptions (strict  => 2);
   if($N3_only_flag) {      
      @programs = ('nu_estimate_np_and_em', 'nu_evaluate');
   }
   else {
      @programs = ('nu_estimate_np_and_em', 'nu_estimate_wm', 'nu_evaluate');
   }
   RegisterPrograms(\@programs);

   # assemble options for N3 method
   if($estimation_type eq 'np') {
     AddDefaultArgs('nu_estimate_np_and_em', 
         [qw(-parzen -log -sharpen), $user_options{'sharpen'}, 0.01, ]);
     foreach $option_name (@option_names) {
       AddDefaultArgs('nu_estimate_np_and_em', ["-$option_name", 
         ((defined $user_options{$option_name})? 
               $user_options{$option_name} :
               $defaults{"np:$option_name:$version_number"})]);
         }
     (defined $user_options{'normalize_field'}) and 
       AddDefaultArgs('nu_estimate_np_and_em','-normalize_field');
   }
   # assemble options for EM method
   elsif($estimation_type eq 'em') {
      AddDefaultArgs('nu_estimate_np_and_em',
                     [qw(-shrink 3 -real -probability default)]);
      foreach $option_name (@option_names) {
         AddDefaultArgs('nu_estimate_np_and_em',
                        ["-$option_name",  
                         (defined $user_options{$option_name})? 
                          $user_options{$option_name} :
                          $defaults{"em:$option_name"}]);
      }
      (defined $user_options{'normalize_field'}) and 
        AddDefaultArgs('nu_estimate_np_and_em','-normalize_field');
   }
   # assemble options for WM method
   else {
      AddDefaultArgs('nu_estimate_wm', 
            ['-spline', "-b_spline -extrapolate "
             ."-distance $user_options{'distance'}"]);
      (defined $user_options{'iterations'}) &&
               AddDefaultArgs('nu_estimate_wm', 
                              ['-N', $user_options{'iterations'}]); 
   }
   # options common to EM and N3
   AddDefaultArgs('nu_estimate_np_and_em',
                  [qw(-auto_mask -nonotify -b_spline 1),
                   '-distance', $user_options{'distance'}]);
   # add common options
   if(defined $user_options{'mask'}) {
      AddDefaultArgs('nu_estimate_np_and_em',['-mask', $user_options{'mask'}]);
      AddDefaultArgs('nu_estimate_wm',['-mask', $user_options{'mask'}])
        unless $N3_only_flag;
   }
   AddDefaultArgs(($N3_only_flag) ? 
        ['nu_estimate_np_and_em', 'nu_evaluate'] :
        ['nu_estimate_np_and_em', 'nu_estimate_wm', 'nu_evaluate'],
                  [(($Verbose) ? '-verbose' : '-quiet'),
                   (($Execute) ? '-execute' : '-noexecute'),
                   (($Clobber) ? '-clobber' :  '-noclobber'),
                   (($KeepTmp) ? '-keeptmp' :  '-nokeeptmp'),
                   '-tmpdir', $TmpDir]);
} 



